from PySide2.QtCore import QFile
from xml.etree import ElementTree
import io
import os
import xml.etree.ElementTree as Et
from PySide2.QtGui import *
from PySide2.QtSvg import *
from PySide2.QtWidgets import *
from PySide2.QtCore import *
import os

def parse_svg(icon_path):
    extension = os.path.splitext(icon_path)[1]
    if extension != '.svg':
        raise Exception(f'{icon_path} 不是一个 SVG 文件。')  # <class'str'> 或 <class'str'> ('.svg' 或 '.
    
    tree = None
    if os.path.exists(icon_path):
        tree = ElementTree.parse(icon_path)
    else:
        file = QFile(icon_path)
        if file.open(QFile.ReadOnly):
            data = file.readAll().data().decode()
            file.close()

            file_obj = io.StringIO(data)
            tree = ElementTree.parse(file_obj)
    return tree

class SvgWidget(QSvgWidget):
    def __init__(self, icon_path, parent=None, color=QColor(7,193,96), size=None):
        super().__init__(parent)
        self.tree = parse_svg(icon_path)
        self.root = self.tree.getroot()
        self.change_color(color)
        self.load(self.get_QByteArray())
        self._size = size
    
    def change_color(self, color):
        c = QColor(color)
        paths = self.root.findall('.//{*}path')
        for path in paths:
            path.set('fill', c.name())
    
    def get_QByteArray(self):
        xmlstr = Et.tostring(self.root, encoding='utf8', method='xml')
        return QByteArray(xmlstr)
    
    def get_QIcon(self):
        # 从 QByteArray 中加载 SVG 数据
        pixmap = QPixmap()
        pixmap.loadFromData(self.get_QByteArray())
        
        # 如果提供了 size 参数，则缩放 QPixmap
        if self._size is not None:
            pixmap = pixmap.scaled(self._size, Qt.KeepAspectRatio, Qt.SmoothTransformation)
        
        # 使用 QPixmap 对象创建 QIcon 对象
        icon = QIcon(pixmap)
        return icon
